#include <iostream>
#include <cmath>
#include <fstream>
#include <vector>
#include <cstdlib>
#include <string>

const double PI = 3.14159265358979323846;

double rho(double w, int n, int k) {
    double r = 1 - 2.0 * w * std::pow(sin(k * PI / (2 * n)), 2);
    return r;
}

int main() {
    // 正确初始化 std::vector<double>
    std::vector<double> w_values = {1.0, 0.5, 1.0/3, 2.0/3};
    int n = 64; // 你可以根据需要调整 n 的值
    int max_k = 63; // 你可以根据需要调整 k 的最大值

    // 为每个 w 值创建单独的数据文件
    for (size_t i = 0; i < w_values.size(); ++i) {
        double w = w_values[i];
        std::string filename = "rho_w" + std::to_string(i) + ".dat";
        std::ofstream dataFile(filename);
        if (!dataFile.is_open()) {
            std::cerr << "无法打开文件 " << filename << std::endl;
            return 1;
        }

        // 写入数据
        for (int k = 1; k <= max_k; k++) {
            double r = rho(w, n, k);
            dataFile << k << " " << r << std::endl;
        }

        dataFile.close();
    }

    // 编写 gnuplot 脚本
    std::ofstream plotScript("plot_rho.gp");
    if (!plotScript.is_open()) {
        std::cerr << "无法打开文件 plot_rho.gp" << std::endl;
        return 1;
    }

    plotScript << "set terminal png size 800,600" << std::endl;
    plotScript << "set output 'rho_plot.png'" << std::endl;
    plotScript << "set title 'rho vs k for different w values'" << std::endl;
    plotScript << "set xlabel 'k'" << std::endl;
    plotScript << "set ylabel 'rho'" << std::endl;
    plotScript << "set xzeroaxis" << std::endl; // 将横坐标轴放在纵坐标轴的0处
    plotScript << "plot ";

    // 添加不同 w 值的数据文件和线型
    for (size_t i = 0; i < w_values.size(); ++i) {
        plotScript << "'rho_w" << i << ".dat' with lines title 'w=" << w_values[i] << "'";
        if (i < w_values.size() - 1) {
            plotScript << ", ";
        }
    }

    plotScript.close();

    // 执行 gnuplot 脚本
    system("gnuplot plot_rho.gp");

    return 0;
}